import numpy as np
import common.Algo_PolicyEvaluation as algo
import common.CommonHelper as helper

# 状态数量
N_States = 13

# 奖励
#    0   1  2  3  4  5  6  7  8  9  10 11 12
R = [-4, 0, 1, 3, 0, 1, 3, 0, 1, 3, 0, 1, 3]

# 转移矩阵
P = np.array(
    [   # 0     1     2     3    4     5      6     7     8     9     10    11    12
        [0.00, 0.56, 0.38, 0.06, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 0
        [0.00, 0.00, 0.00, 0.00, 0.56, 0.38, 0.06, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 1
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.50, 0.42, 0.08, 0.00, 0.00, 0.00], # 2
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.40, 0.50, 0.10], # 3
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 4
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 5
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 6
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 7
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 8
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 9
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 10
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 11
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00], # 12
    ]
)

# 简单模型
class DataModel(object):
    def __init__(self):
        self.P = P              # 状态转移矩阵
        self.R = R              # 奖励
        self.nS = N_States      # 状态数量

if __name__=="__main__":
    gamma = 1.0
    dataModel = DataModel()
    V = algo.solve_matrix(dataModel, gamma)
    V = np.round(V, 2)
    print("状态价值函数：")
    helper.print_seperator_line(helper.SeperatorLines.middle)
    print("v0 =", V[0])
    helper.print_seperator_line(helper.SeperatorLines.short)
    print(str.format("v1={0},\tv2={1},\tv3={2}", V[1],V[2],V[3]))
    helper.print_seperator_line(helper.SeperatorLines.short)
    print(str.format("v4={0},\t\tv5={1},\t\tv6={2}", V[4],V[5],V[6]))
    print(str.format("v7={0},\t\tv8={1},\t\tv9={2}", V[7],V[8],V[9]))
    print(str.format("v10={0},\tv11={1},\tv12={2}", V[10],V[11],V[12]))
